package org.example.test.controller;

import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;
import java.time.Duration;

/**
 * 流式聊天API控制器
 * 使用Spring WebFlux实现服务器推送事件（SSE）
 */
@RestController
public class ChatController {

    // 自动注入Ollama聊天客户端（需提前配置Bean）
    @Autowired
    private OllamaChatClient chatClient;



    /**
     * 流式聊天接口
     * @param message 用户输入消息
     * @return ServerSentEvent流（text/event-stream格式）
     */
    @GetMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> streamResponse(@RequestParam String message) {
        // 创建响应流处理器（带背压控制的单播Sink）
        // - unicast(): 单订阅者模式
        // - onBackpressureBuffer(): 启用背压缓冲
        Sinks.Many<ServerSentEvent<String>> sink = Sinks.many().unicast().onBackpressureBuffer();

        // 构建Ollama模型参数（完整参数说明见下文）
        OllamaOptions options = OllamaOptions.create()
                .withTemperature(0.7f)  // 随机性控制（0-1，值越大输出越随机）
                .withTopK(50)           // 候选词保留数量（控制输出多样性）
                .withNumPredict(200);   // 最大生成token数（防止无限生成）

        // 发起流式AI请求
        // 参数说明：
        // - Prompt: 包含用户消息和模型参数的请求对象
        // - chatClient.stream():  返回Flux<ChatResponse>的响应流
        Flux<ChatResponse> aiResponse = chatClient.stream(
                new Prompt(message, options)
        );

        // 处理AI响应流
        aiResponse
                // 设置30秒超时（超过自动终止流）
                .timeout(Duration.ofSeconds(30))
                // 处理每个响应块
                .doOnNext(response -> {
                    // 解析响应内容（根据模型返回粒度可能是词/句子）
                    String content = response.getResult().getOutput().getContent();

                    // 将内容拆分为token级流（按单词边界分割）
                    for (String token : content.split("(?<=\\b)"))  {
                        // 构建SSE事件（默认事件类型为"message"）
                        sink.tryEmitNext(ServerSentEvent.builder(token).build());
                    }
                })
                // 处理流完成
                .doOnComplete(() -> {
                    // 发送结束标记事件
                    sink.tryEmitNext(ServerSentEvent.<String>builder()
                            .event("[END]")  // 自定义事件类型
                            .data("")        // 空数据
                            .build());
                    // 关闭流
                    sink.tryEmitComplete();
                })
                // 处理错误
                .doOnError(e -> {
                    // 发送错误事件
                    sink.tryEmitNext(ServerSentEvent.<String>builder()
                            .event("[ERROR]")    // 自定义错误事件类型
                            .data(e.getMessage())  // 携带错误信息
                            .build());
                    // 关闭流
                    sink.tryEmitComplete();
                })
                // 启动订阅（无订阅不执行）
                .subscribe();

        // 返回响应流，并处理客户端断开连接
        return sink.asFlux()
                // 客户端取消订阅时的处理
                .doOnCancel(() -> {
                    // 立即关闭流（释放资源）
                    sink.tryEmitComplete();
                });
    }
}
